import os
import re
import threading
import tempfile
import time
import win32com.client as win32
import pythoncom
import tkinter as tk
from tkinter import scrolledtext

# 第三方库
from docx import Document as DocxDocument
from pptx import Presentation as PptxPresentation
import openpyxl


class DocumentSearchApp:
    def __init__(self, root):
        self.root = root
        self.root.title("办公文档关键字搜索")

        # 默认关键字
        self.default_keywords = "秘密,机密,绝密"

        # 创建界面
        self.create_widgets()

    def create_widgets(self):
        # 关键字输入框
        input_frame = tk.Frame(self.root)
        input_frame.pack(pady=10)

        tk.Label(input_frame, text="关键字（逗号分隔）：").pack(side=tk.LEFT)
        self.keyword_entry = tk.Entry(input_frame, width=50)
        self.keyword_entry.pack(side=tk.LEFT, padx=5)
        self.keyword_entry.insert(0, self.default_keywords)
        self.search_button = tk.Button(input_frame, text="检索", command=self.start_search_thread)
        self.search_button.pack(side=tk.LEFT)

        # 日志输出
        tk.Label(self.root, text="日志信息：").pack(anchor="w")
        self.log_text = scrolledtext.ScrolledText(self.root, height=10, width=80)
        self.log_text.pack(padx=10, pady=5)

        # 搜索结果
        tk.Label(self.root, text="搜索结果：").pack(anchor="w")
        self.result_text = scrolledtext.ScrolledText(self.root, height=15, width=80)
        self.result_text.pack(padx=10, pady=5)

        # 版权信息
        copyright_label = tk.Label(self.root, text="文山州公安局刑侦支队 张泽山@2025")
        copyright_label.pack(pady=5)

    def log(self, message):
        self.root.after(0, lambda: self.log_text.insert(tk.END, message + "\n"))
        self.root.after(0, lambda: self.log_text.see(tk.END))

    def result(self, message):
        self.root.after(0, lambda: self.result_text.insert(tk.END, message + "\n"))
        self.root.after(0, lambda: self.result_text.see(tk.END))

    def extract_text_from_docx(self, file_path):
        try:
            doc = DocxDocument(file_path)
            full_text = []
            for para in doc.paragraphs:
                full_text.append(para.text)
            for table in doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        full_text.append(cell.text)
            return "\n".join(full_text)
        except Exception as e:
            self.log(f"无法读取 .docx 文件 {file_path}: {str(e)}")
            return ""

    def extract_text_from_xlsx(self, file_path):
        try:
            workbook = openpyxl.load_workbook(file_path)
            full_text = []
            for sheet_name in workbook.sheetnames:
                sheet = workbook[sheet_name]
                for row in sheet.iter_rows():
                    for cell in row:
                        if cell.value:
                            full_text.append(str(cell.value))
            return "\n".join(full_text)
        except Exception as e:
            self.log(f"无法读取 .xlsx 文件 {file_path}: {str(e)}")
            return ""

    def extract_text_from_pptx(self, file_path):
        try:
            pres = PptxPresentation(file_path)
            full_text = []
            for slide in pres.slides:
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        full_text.append(shape.text)
            return "\n".join(full_text)
        except Exception as e:
            self.log(f"无法读取 .pptx 文件 {file_path}: {str(e)}")
            return ""

    def extract_text_from_doc(self, file_path):
        try:
            word = win32.gencache.EnsureDispatch('Word.Application')
            word.Visible = False
            doc = word.Documents.Open(file_path)
            content = doc.Content.Text
            doc.Close(False)
            word.Quit()
            return content
        except Exception as e:
            self.log(f"无法读取 .doc 文件 {file_path}: {str(e)}")
            return ""

    def extract_text_from_xls(self, file_path):
        return self.extract_text_with_excel(file_path)

    def extract_text_from_ppt(self, file_path):
        return self.extract_text_with_powerpoint(file_path)

    def extract_text_with_excel(self, file_path):
        try:
            excel = win32.gencache.EnsureDispatch('Excel.Application')
            excel.Visible = False
            workbook = excel.Workbooks.Open(file_path)
            full_text = []
            for sheet in workbook.Sheets:
                for row in range(1, sheet.UsedRange.Rows.Count + 1):
                    for col in range(1, sheet.UsedRange.Columns.Count + 1):
                        cell_value = sheet.Cells(row, col).Value
                        if cell_value:
                            full_text.append(str(cell_value))
            workbook.Close(False)
            excel.Quit()
            return "\n".join(full_text)
        except Exception as e:
            self.log(f"无法读取 .xls 文件 {file_path}: {str(e)}")
            return ""

    def extract_text_with_powerpoint(self, file_path):
        try:
            powerpoint = win32.gencache.EnsureDispatch('PowerPoint.Application')
            powerpoint.Visible = False
            presentation = powerpoint.Presentations.Open(file_path)
            full_text = []
            for slide in presentation.Slides:
                for shape in slide.Shapes:
                    if hasattr(shape, "TextFrame") and shape.HasTextFrame:
                        if shape.TextFrame.HasText:
                            full_text.append(shape.TextFrame.TextRange.Text)
            presentation.Close(False)
            powerpoint.Quit()
            return "\n".join(full_text)
        except Exception as e:
            self.log(f"无法读取 .ppt 文件 {file_path}: {str(e)}")
            return ""

    def search_file(self, file_path, keywords):
        ext = os.path.splitext(file_path)[1].lower()
        content = ""

        try:
            if ext == ".docx":
                content = self.extract_text_from_docx(file_path)
            elif ext == ".xlsx":
                content = self.extract_text_from_xlsx(file_path)
            elif ext == ".pptx":
                content = self.extract_text_from_pptx(file_path)
            elif ext == ".doc":
                content = self.extract_text_from_doc(file_path)
            elif ext == ".xls":
                content = self.extract_text_from_xls(file_path)
            elif ext == ".ppt":
                content = self.extract_text_from_ppt(file_path)
            else:
                return []

            results = []
            for keyword in keywords:
                matches = re.finditer(re.escape(keyword), content, re.IGNORECASE)
                for match in matches:
                    context = content[max(0, match.start() - 25):match.end() + 25]
                    results.append((keyword, context))
            return results
        except Exception as e:
            self.log(f"处理文件失败 {file_path}: {str(e)}")
            return []

    def safe_walk(self, directory):
        """ 安全版本的 os.walk，跳过权限错误 """
        try:
            for root, dirs, files in os.walk(directory):
                yield root, dirs, files
        except Exception as e:
            self.log(f"[错误] 访问目录失败: {directory} - {str(e)}")

    def get_local_drives(self):
        """ 获取本地磁盘（如 C:\\, D:\\ 等存在的盘符） """
        drives = []
        for drive in range(ord("A"), ord("Z") + 1):
            drive_name = f"{chr(drive)}:\\"
            if os.path.exists(drive_name):
                drives.append(drive_name)
        return drives

    def start_search_thread(self):
        self.search_button.config(state=tk.DISABLED)
        self.log_text.delete(1.0, tk.END)
        self.result_text.delete(1.0, tk.END)
        thread = threading.Thread(target=self.start_search)
        thread.start()

    def start_search(self):
        try:
            pythoncom.CoInitialize()  # 初始化 COM 库
            keywords = [k.strip() for k in self.keyword_entry.get().split(",") if k.strip()]
            if not keywords:
                self.log("请输入至少一个关键字")
                self.search_button.config(state=tk.NORMAL)
                return

            drives = self.get_local_drives()
            self.log(f"开始扫描以下磁盘: {', '.join(drives)}")

            for drive in drives:
                self.log(f"开始扫描磁盘 {drive}")
                for root, dirs, files in self.safe_walk(drive):
                    self.log(f"进入目录: {root}")
                    for file in files:
                        try:
                            file_path = os.path.join(root, file)
                            ext = os.path.splitext(file)[1].lower()
                            if ext in (".docx", ".xlsx", ".pptx", ".doc", ".xls", ".ppt"):
                                self.log(f"正在搜索文件: {file_path}")
                                results = self.search_file(file_path, keywords)
                                if results:
                                    self.result(f"找到匹配项在: {file_path}")
                                    for keyword, context in results:
                                        self.result(f"关键字: {keyword}")
                                        self.result(f"上下文: {context}")
                                        self.result("-" * 50)
                        except Exception as e:
                            self.log(f"跳过文件: {file} - {str(e)}")
                            continue

            self.log("搜索完成")
        finally:
            self.search_button.config(state=tk.NORMAL)


if __name__ == "__main__":
    root = tk.Tk()
    app = DocumentSearchApp(root)
    root.mainloop()



